import torch
import torch.nn.functional as F
from tqdm import tqdm
import torch.nn as nn
from dice_loss import dice_coeff

def eval_net(net, loader, device, out_channel=1, global_step=0, writer=None):
    """Evaluation without the densecrf with the dice coefficient"""
    net.eval()
    mask_type = torch.float32 if out_channel == 1 else torch.long
    n_val = len(loader)  # the number of batch
    n_val2 = len(loader)
    tot = 0
    val_loss = 0
    criterion = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([10], device=device))
    flag = 0 # 如果添加了一张有肿瘤的图像，就可以结束了。
    with tqdm(total=n_val, desc='Validation round', unit='batch', leave=False) as pbar:
        for imgs, true_masks in loader:
            batch, channe, D, W, H = imgs.shape
            # imgs = imgs.view((batch, 1, 1, W, H))
            imgs = imgs.to(device=device, dtype=torch.float32)
            true_masks = true_masks.to(device=device, dtype=mask_type)

            with torch.no_grad():
                mask_pred = net(imgs)
                # mask_pred = mask_pred.view((batch, W, H))
                mask_pred = mask_pred.squeeze(1) # 去掉通道那个维度
            if out_channel > 1:
                tot += F.cross_entropy(mask_pred, true_masks).item()
            else:
                val_loss += criterion(mask_pred, true_masks) # 计算验证集的损失
                
                pred = torch.sigmoid(mask_pred)
                if true_masks.sum() != 0:
                    tot += dice_coeff(pred, true_masks).item()
                    if flag == 0:
                        # 说明还没添加图像
                        imgs = imgs.squeeze(2) # 去掉深度维度
                        masks_pred = mask_pred.squeeze(2) # 去掉深度维度
                        writer.add_images('images', imgs, global_step)
                        
                        writer.add_images('masks/true', true_masks, global_step)
                        writer.add_images('masks/pred', (masks_pred > 0.5), global_step)
                        flag += 1
                else :
                    n_val -= 1

            pbar.update()

    net.train()
    dice_loss = tot / n_val
    val_loss = val_loss / n_val2

    return dice_loss, val_loss  
